import os
import shutil

from mAP.mAP import get_mAP_by_dir


resolutions = [240, 360, 480, 640, 720, 960, 1080, 1280]
currentPath = os.path.dirname(os.path.realpath(__file__))


def predict2gt():
    for scene in range(1, 9):
        predict_dir = "datasets/AuAir/scene" + str(scene) + "/predict1280"
        gt_dir = "datasets/AuAir/scene" + str(scene) + "/anno1280"
        if not os.path.exists(gt_dir):
            os.mkdir(gt_dir)
        else:
            shutil.rmtree(gt_dir)
            os.mkdir(gt_dir)
        file_names = os.listdir(predict_dir)
        for file_name in file_names:
            with open(os.path.join(predict_dir, file_name), 'r', ) as file:
                info = file.readlines()
            with open(os.path.join(gt_dir, file_name), 'a', ) as file:
                for item in info:
                    item = item.replace("\n", '').split(' ')
                    line = "{} {} {} {} {}\n".format(item[0], item[2], item[3], item[4], item[5])
                    file.write(line)


def calc_mAP():
    for scene in range(1, 9):
        print("场景" + str(scene) + "下的检测结果")
        gt_path = os.path.join(currentPath, "datasets/AuAir/scene" + str(scene) + "/anno1280")
        for resolution in resolutions:
            det_path = os.path.join(currentPath, "datasets/AuAir/scene" + str(scene) + "/predict" + str(resolution))
            result = get_mAP_by_dir(det_path, gt_path)
            print('分辨率为{}下的mAP值为: {} \n'.format(resolution, result))


predict2gt()
calc_mAP()
